/* SPDX-FileCopyrightText: 2021 John Scott <jscott@posteo.net>
 * SPDX-License-Identifier: GPL-3.0-or-later */

/* We do not support the obsolete extension where an
 * omitted directory name is interpreted as the current
 * working directory. In $PATH = "/usr::/bin:", the lack
 * of a path in the middle or one at the end is simply ignored. */

#define _XOPEN_SOURCE 700
#include <assert.h>
#include <errno.h>
#include <limits.h>
#include <locale.h>
#include <pthread.h>
#include <search.h>
#include <semaphore.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>
#if _POSIX_ADVISORY_INFO > -1
#include <sys/mman.h>
#endif

/* A NULL-terminated list of directories in PATH. */
static char **list;

/* We need a way to tell which threads are running at any given time
 * so we know which ones we can send cancellation requests to.
 * This is a boolean list indicating whether a given thread is
 * running or not. thread_is_running[0] will correspond to the
 * first child we create, thread_is_running[1] to the second,
 * and so on. */
static bool *thread_is_running;

/* This mutex not only protects thread_is_running from data races,
 * but also has the effect that, if we lock it in the main thread,
 * child threads will be waiting for this to become unlocked just
 * before they set thread_is_running[j] to false and bail out.
 * This means that locking this mutex in the main thread will
 * freeze which threads are running or not, so that we may safely
 * send cancellation requests to those that are. */
static pthread_mutex_t thread_is_running_guard = PTHREAD_MUTEX_INITIALIZER;

/* We need indices into thread_is_running. This integer is
 * an identifier for the thread which may be used as an index.
 * We increment it, starting at 0, for every thread we create. */
static int *seq_thread_id;

/* We need to give every child thread a chance to read seq_thread_id
 * before clobbering it. This semaphore ensures that we do that; the
 * main thread will wait for it to be unlocked by the child, and the
 * child will unlock it after it has copied the value. */
static sem_t seq_thread_id_guard;

/* This is a list of groups we are in. Its lifetime is managed by main(). */
static gid_t *groups;
static int groupcount;

static void *reallocarray(void *p, size_t m, size_t n) {
	if(n && m > SIZE_MAX / n) {
		errno = ENOMEM;
		return NULL;
	}
	return realloc(p, m * n);
}

static int giddiff(const void *a, const void *b) {
	gid_t gid_a = *(const gid_t*)a;
	gid_t gid_b = *(const gid_t*)b;
	/* We do not simply return gid_a - gid_b, because that
	 * bears the risk of overflow if gid_t is a signed type. */
	if(gid_a > gid_b) {
		return 1;
	} else if(gid_a < gid_b) {
		return -1;
	} else {
		return 0;
	}
}

static void stop_running(void *i) {
	int k = pthread_mutex_lock(&thread_is_running_guard);
	if(k) {
		char errstr[NL_TEXTMAX];
		if(strerror_r(k, errstr, sizeof(errstr))) {
			abort();
		}
		fprintf(stderr, "Failed to lock mutex: %s\n", errstr);
		abort();
	}
	assert(thread_is_running[*(int*)i]);
	thread_is_running[*(int*)i] = false;
	k = pthread_mutex_unlock(&thread_is_running_guard);
	if(k) {
		abort();
	}
}

static bool file_is_executable(const char filename[restrict static 1]) {
	struct stat st;
	if(stat(filename, &st) == -1 || !S_ISREG(st.st_mode)) {
		return false;
	}

	if(st.st_uid == geteuid()) {
		return (st.st_mode & S_IXUSR) ? true : false;
	}

	if(bsearch(&st.st_gid, groups, groupcount, sizeof(*groups), giddiff)) {
		return (st.st_mode & S_IXGRP) ? true : false;
	}

	return (st.st_mode & S_IXOTH) ? true : false;
}

/* This is the main function that corresponds to our child thread.
 * It tries to find the program called filename in our list of
 * directories in PATH and to return a dynamically-allocated
 * string to the full pathname, or NULL if it can't do that. */
static void *find_program(void *filename) {
	int my_seq_thread_id = *seq_thread_id;
	if(sem_post(&seq_thread_id_guard) == -1) {
		abort();
	}

	int k = pthread_mutex_lock(&thread_is_running_guard);
	if(k) {
		char errstr[NL_TEXTMAX];
		if(strerror_r(k, errstr, sizeof(errstr))) {
			abort();
		}
		fprintf(stderr, "Failed to lock mutex: %s\n", errstr);
		abort();
	}
	assert(!thread_is_running[my_seq_thread_id]);
	thread_is_running[my_seq_thread_id] = true;
	k = pthread_mutex_unlock(&thread_is_running_guard);
	if(k) {
		abort();
	}

	/* These have to be declared outside of our cleanup handler calls. */
	char *pathname;
	bool pathname_found = false;
	/* We hit no cancellation points between setting thread_is_running[j]
	 * to true and pushing this handler, since pthread_mutex_unlock is not one. */
	pthread_cleanup_push(stop_running, &my_seq_thread_id);

	if(strchr(filename, '/')) {
		/* We were either given an absolute pathname, or a relative
		 * path that must be evaluated with respect to the current
		 * working directory, which must not use prefixes from PATH. */
		if(file_is_executable(filename)) {
			/* The caller expects that the string we return is dynamically allocated. */
			pathname = strdup(filename);
			if(!pathname) {
				perror("Failed to duplicate string");
			}
			pthread_exit(pathname);
		} else {
			pthread_exit(NULL);
		}
	}

	/* It's not an absolute pathname; try prefixing the string
	 * with strings from PATH and see what sticks. */
	const size_t filenamelen = strlen(filename);
	for(char **directory = list; *directory; directory++) {
		const size_t directorylen = strlen(*directory);
		if(filenamelen > SIZE_MAX - directorylen || filenamelen + directorylen > SIZE_MAX - 2) {
			continue;
		}
		pathname = malloc(directorylen + filenamelen + 2); /* one extra byte for a /, one for the NUL */
		if(!pathname) {
			perror("Failed to allocate memory for pathname");
			pthread_exit(NULL);
		}
		pthread_cleanup_push(free, pathname);

		char *end = (char*)memcpy(pathname, *directory, directorylen) + directorylen;
		if(end[-1] != '/') {
			*end++ = '/';
		}
		memcpy(end, filename, filenamelen + 1);

		pathname_found = file_is_executable(pathname);
		pthread_cleanup_pop(!pathname_found); /* free(pathname)? */
		if(pathname_found) {
			break;
		}
	}

	pthread_cleanup_pop(true); /* stop_running(&my_seq_thread_id) */
	pthread_exit(pathname_found ? pathname : NULL);
}

int main(int argc, char *argv[]) {
	if(!setlocale(LC_ALL, "")) {
		fputs("Failed to enable default locale\n", stderr);
		exit(EXIT_FAILURE);
	}

	int opt;
	while((opt = getopt(argc, argv, "")) != -1) {
		if(opt == '?') {
			exit(EXIT_FAILURE);
		}
	}
	argc -= optind;
	argv += optind;
	if(!argc) {
		exit(EXIT_SUCCESS);
	}

	const char *const envpath = getenv("PATH");
	char *path;
	size_t l;
	if(!envpath || !envpath[0]) {
		l = confstr(_CS_PATH, NULL, 0);
		if(!l) {
			fputs("Failed to obtain value of PATH\n", stderr);
			exit(EXIT_FAILURE);
		}

		path = aligned_alloc(sysconf(_SC_PAGESIZE), l);
		if(!path) {
			perror("Failed to allocate memory for PATH");
			exit(EXIT_FAILURE);
		}
		confstr(_CS_PATH, path, l);
	} else {
		l = strlen(envpath) + 1;
		path = aligned_alloc(sysconf(_SC_PAGESIZE), l);
		if(!path) {
			perror("Failed to duplicate string");
			exit(EXIT_FAILURE);
		}
		memcpy(path, envpath, l);
	}
#if _POSIX_ADVISORY_INFO > -1
	int p = posix_madvise(path, l, POSIX_MADV_SEQUENTIAL);
	if(p && sysconf(_SC_ADVISORY_INFO) != -1) {
		fprintf(stderr, "Failed to advise the system on memory usage: %s\n", strerror(p));
	}
#endif

	/* The maximum number of directories in PATH is one plus
	 * the number of colons, where multiple consecutive colons
	 * can be treated as a single one. */
	size_t numdirs = 1;
	assert(path[0]);
	for(size_t i = 1; path[i]; i++) {
		/* In case of a set of multiple consecutive colons,
		 * only count the last one. */
		if(path[i] == ':' && path[i+1] && path[i+1] != ':') {
			numdirs++;
		}
	}

	assert(numdirs < SIZE_MAX);
	list = reallocarray(NULL, numdirs + 1, sizeof(*list));
	if(!list) {
		perror("Failed to allocate memory for directory list");
		goto endpath;
	}

	char *tok = NULL;
	size_t n = 0;
	do {
		tok = strtok(tok ? NULL : path, ":");
		assert(n <= numdirs);
		list[n++] = tok;
	} while(tok);

	pthread_t *ids = reallocarray(NULL, argc, sizeof(*ids));
	if(!ids) {
		perror("Failed to allocate memory for thread list");
		goto endlist;
	}

	if((groupcount = getgroups(0, groups)) == -1) {
		perror("Failed to get number of groups");
		goto endids;
	}
#pragma GCC diagnostic ignored "-Wsign-compare"
	if(groupcount == INT_MAX || groupcount == SIZE_MAX) {
#pragma GCC diagnostic pop
		fprintf(stderr, "Failed to create group list: %s\n", strerror(EOVERFLOW));
		goto endids;
	}

	/* We might need an extra member for the effective group ID. */
	groups = reallocarray(NULL, groupcount + 1, sizeof(*groups));
	if(!groups) {
		perror("Failed to allocate memory for group list");
		goto endids;
	}
	/* It's possible that in a TOCTTOU sort of way, the number of
	 * groups we're in now is fewer than the number we were in before,
	 * hence the reassignment to groupcount. */
	if((groupcount = getgroups(groupcount, groups)) == -1) {
		perror("Failed to populate group list");
		goto endgroups;
	}

	/* The group list may not include the effective group ID. */
	if(!lfind(&(gid_t){getegid()}, groups, &(size_t){groupcount}, sizeof(*groups), giddiff)) {
		groups[groupcount++] = getegid();
	}
	qsort(groups, groupcount, sizeof(*groups), giddiff);

	if(sem_init(&seq_thread_id_guard, false, 0U) == -1) {
		perror("Failed to initialize semaphore");
		goto endgroups;
	}

	thread_is_running = calloc(argc, sizeof(*thread_is_running));
	if(!thread_is_running) {
		perror("Failed to allocate memory for running thread list");
		goto endseq_thread_id_guard;
	}

	void **retval = reallocarray(NULL, argc, sizeof(*retval));
	if(!retval) {
		perror("Failed to allocate memory for thread return values");
		goto endthread_is_running;
	}

	int i;
	seq_thread_id = &i;
	for(i = 0; i < argc; i++) {
		int k;
tryagain:
		k = pthread_create(ids + i, NULL, find_program, argv[i]);
		switch(k) {
		case 0:
			break;
		case EAGAIN:
			if(sched_yield() == -1) {
				perror("Failed to yield");
			}
			goto tryagain;
		default:
			fprintf(stderr, "Failed to create thread: %s\n", strerror(k));
			k = pthread_mutex_lock(&thread_is_running_guard);
			if(k) {
				fprintf(stderr, "Failed to lock mutex: %s\n", strerror(k));
				abort();
			}

			for(int j = 0; j < i; j++) {
				if(thread_is_running[j]) {
					k = pthread_cancel(ids[j]);
					if(k) {
						fprintf(stderr, "Failed to cancel thread: %s\n", strerror(k));
						abort();
					}
				}
			}
			k = pthread_mutex_unlock(&thread_is_running_guard);
			if(k) {
				fprintf(stderr, "Failed to unlock mutex: %s\n", strerror(k));
				abort();
			}

			void *threadreturn;
			for(int j = 0; j < i; j++) {
				k = pthread_join(ids[j], &threadreturn);
				if(k) {
					fprintf(stderr, "Failed to join with thread: %s\n", strerror(k));
					abort();
				}
				if(threadreturn != PTHREAD_CANCELED) {
					free(threadreturn);
				}
			}
			goto endthread_is_running_guard;
		}

		if(sem_wait(&seq_thread_id_guard) == -1) {
			perror("Failed to lock semaphore");
			abort();
		}
	}

	for(int j = 0; j < argc; j++) {
		int k = pthread_join(ids[j], retval + j);
		if(k) {
			fprintf(stderr, "Failed to join with thread: %s\n", strerror(k));
			abort();
		}
	}

	bool all_found = true;
	for(int j = 0; j < argc; j++) {
		if(retval[j]) {
			if(puts(retval[j]) == EOF) {
				perror("Failed to print filename");
				all_found = false;
			}
			free(retval[j]);
		} else {
			all_found = false;
		}
	}

	int k = pthread_mutex_destroy(&thread_is_running_guard);
	if(k) {
		fprintf(stderr, "Failed to destroy mutex: %s\n", strerror(k));
		abort();
	}
	free(retval);
	if(sem_destroy(&seq_thread_id_guard) == -1) {
		perror("Failed to destroy semaphore");
		abort();
	}
	free(groups);
	free(ids);
	free(list);
	free(path);
	exit(all_found ? EXIT_SUCCESS : EXIT_FAILURE);

endthread_is_running_guard:
	k = pthread_mutex_destroy(&thread_is_running_guard);
	if(k) {
		fprintf(stderr, "Failed to destroy mutex: %s\n", strerror(k));
		abort();
	}
	free(retval);
endthread_is_running:
	free(thread_is_running);
endseq_thread_id_guard:
	if(sem_destroy(&seq_thread_id_guard) == -1) {
		perror("Failed to destroy semaphore");
		abort();
	}
endgroups:
	free(groups);
endids:
	free(ids);
endlist:
	free(list);
endpath:
	free(path);
	exit(EXIT_FAILURE);
}
